package com.math.aistudykbr.controller;

import com.math.aistudykbr.config.SSEConnectManager;
import com.math.aistudykbr.utils.UserContext;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;

@RestController
@Slf4j
@RequestMapping("/api/sse")
@RequiredArgsConstructor
public class SSEController {

//    private static final ConcurrentHashMap<String, SseEmitter> CONNECTION_POOL = new ConcurrentHashMap<>();

    private final SSEConnectManager sseConnectManager;

    private static final long TIMEOUT = 3_600_000L; // 1小时超时

    @GetMapping(value = "/connect", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter connect() {
        // 1. 获取当前用户ID
        Long userId = UserContext.getUser();

        // 2. 创建新连接（超时一个小时）
        SseEmitter emitter = new SseEmitter(TIMEOUT);
        log.info("用户 {} 建立SSE连接", userId);

        // 3. 加入连接池，保证原子性
        sseConnectManager.addConnection(userId.toString(), emitter);

        // 4. 发送连接成功消息
        try {
            emitter.send(SseEmitter.event()
                    .data(ResponseEntity.ok()
                            .contentType(MediaType.TEXT_PLAIN)
                            .body("SSE连接已建立"))
                    .id("connect-" + System.currentTimeMillis())
            );
        } catch (IOException e) {
            emitter.completeWithError(e);
        }

        // 5. 设置清理回调，当连接关闭时，从连接池中移除
        emitter.onCompletion(() -> {
            sseConnectManager.removeConnection(userId.toString(), emitter);
            log.info("用户 {} 连接正常关闭", userId);
        });
        emitter.onTimeout(() -> {
            sseConnectManager.removeConnection(userId.toString(), emitter);
            log.warn("用户 {} 连接超时", userId);
        });

        return emitter;
    }
}